/*
 *  wbos - parser of pe files
 *
 *  Copyright (C) 2011 by Volodymyr M. Shcherbyna <volodymyr@shcherbyna.com>
 *
 *      This file is part of wbos.
 *
 *  wbos is free software: you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation, either version 3 of the License, or
 *  (at your option) any later version.
 *
 *  wbos is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License
 *  along with wbos.  If not, see <http://www.gnu.org/licenses/>.
*/

#include "peparser.h"

/**   C'tor of the object

    * @return initialized object
*/
CPEParser::CPEParser(char * szFileName) : m_szFileName(szFileName)
{    
}

/**   Calculating the image base of pe binary

    * @return long
*/
unsigned long CPEParser::getImageBase(void)
{
	FILE *f = NULL;

	fopen_s(&f, m_szFileName, "r");

    if (f == NULL)
    {
        /// file does not exist
        std::cout << "The file '" << m_szFileName << "' does not exist or cannot be opened" << std::endl;

        return 0;
    }

	fseek(f, 0, SEEK_CUR);

    IMAGE_DOS_HEADER ImageNtHeader;

	fread(&ImageNtHeader, sizeof(unsigned char), sizeof(ImageNtHeader), f);

    if (ImageNtHeader.e_magic != IMAGE_DOS_SIGNATURE)
    {
        std::cout << "File '" << m_szFileName << "' is not a valid PE file" << std::endl;

        fclose(f);

        return 0;
    }

	fseek(f, ImageNtHeader.e_lfanew, SEEK_SET);

	PARTIAL_NT_HEADER ntHeaders;

    if (fread(&ntHeaders, sizeof(unsigned char), sizeof(ntHeaders), f) != sizeof(ntHeaders))
    {
        std::cout << "File '" << m_szFileName << "' is not a valid PE file. Unable to read Image File Header" << std::endl;

		fclose(f);

        return 0;
    }

    if (ntHeaders.FileHeader.Machine != IMAGE_FILE_MACHINE_I386 && IMAGE_FILE_MACHINE_AMD64 != ntHeaders.FileHeader.Machine)
    {
        std::cout << "File '" << m_szFileName << "' is not a valid PE file. Unable to detect architecture" << std::endl;

        fclose(f);

        return 0;
    }

    unsigned long long nRetVal = 0;
    
    if (ntHeaders.FileHeader.Machine == IMAGE_FILE_MACHINE_I386)
    {
        IMAGE_OPTIONAL_HEADER OptionalHeader32;

		if (fread(&OptionalHeader32, sizeof(unsigned char), sizeof(OptionalHeader32), f) != sizeof(OptionalHeader32))
        {
            std::cout << "File '" << m_szFileName << "' is not a valid PE file. Unable to read Image File Header" << std::endl;

            fclose(f);

            return 0;
        }

        nRetVal = OptionalHeader32.ImageBase;
    }
    else
    {
        IMAGE_OPTIONAL_HEADER64 OptionalHeader64;

		if (fread(&OptionalHeader64, sizeof(unsigned char), sizeof(OptionalHeader64), f) != sizeof(OptionalHeader64))
        {
            std::cout << "File '" << m_szFileName << "' is not a valid PE file. Unable to read Image File Header" << std::endl;
            fclose(f);

            return 0;
        }

        nRetVal = OptionalHeader64.ImageBase;
    }

    fclose(f);

    return (unsigned long)nRetVal;
}